import db
import bson

try:
    import json
except ImportError:
    import simplejson as json
    
J = json.JSONEncoder()

class DB(db.DB):
    def execute(self, input):
        return str(self.c.code.insert({'input':input}, safe=True))

    def get(self, id):
        v = list(self.c.code.find({'_id':bson.objectid.ObjectId(id)}))
        if len(v) == 0:
            # return an error message
            return {'status':'invalid_id'}
        else:
            if v[0].has_key('output'):
                # then done
                return {'status':'done', 'output':v[0]['output']}
            else:
                # still running
                return {'status':'wait'}

    def get_json(self, id):
        return J.encode(self.get(id))
    
    def get_unevaluated_cells(self):
        return self.c.code.find({'output':{'$exists': False}})
    
    def get_evaluated_cells(self):
        import pymongo
        return self.c.code.find({'output':{'$exists': True}}).sort('_id', direction=pymongo.DESCENDING)
    
    def set_output(self, id, output):
        self.c.code.update({'_id':id}, {'$set':{'output':output}})
    
